import pyrealsense2 as rs
import numpy as np
import cv2
from ultralytics import YOLO
import panda_py
from panda_py import libfranka



# Initialize Panda robot and gripper
panda = panda_py.Panda("172.16.0.2")
hand = libfranka.Gripper("172.16.0.2")

# Step 1: Move the robot arm to the top
panda.move_to_start()

# Load YOLOv8 model
model = YOLO('/home/robot1/JuliusWorkplace/Tongji_Intelligent_Systems_Lab/Teaching_Scripts/RealSense/yolov8n.pt')

# Configure RealSense pipeline
pipeline = rs.pipeline()
config = rs.config()
config.enable_stream(rs.stream.color, 640, 480, rs.format.bgr8, 30)
config.enable_stream(rs.stream.depth, 640, 480, rs.format.z16, 30)
profile = pipeline.start(config)



try:
    while True:
        # Capture frames
        frames = pipeline.wait_for_frames()
        color_frame = frames.get_color_frame()
        depth_frame = frames.get_depth_frame()
        if not color_frame or not depth_frame:
            continue

        # Convert frames to numpy arrays
        color_image = np.asanyarray(color_frame.get_data())
        depth_image = np.asanyarray(depth_frame.get_data())

        

        # Perform object detection
        results = model(color_image)

        
        for r in results:
            boxes = r.boxes
            for box in boxes:
                x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
                conf = float(box.conf[0])
                cls = int(box.cls[0])
                cx, cy = (x1 + x2) // 2, (y1 + y2) // 2
                depth = depth_image[cy, cx]
                depth_m = depth / 1000.0

                # Step 2: Move the arm straight above the object
                myPose = panda.get_pose()
                myPose[0, 3] = cx  # Update X position
                myPose[1, 3] = cy  # Update Y position
                q = panda_py.ik(myPose)
                panda.move_to_joint_position(q)




                # Step 3: Move down until the plane to grasp the object
                myPose[2, 3] -= depth_m  # Update Z position
                q = panda_py.ik(myPose)
                panda.move_to_joint_position(q)
                hand.grasp(0.025, 50.0, 0.2)
                break
        break
finally:
    pipeline.stop()
    cv2.destroyAllWindows()